#include "Disk.h"

Disk::Disk(wstring wDiskName, EDeviceType m_eDevType, EIOType	eIoType) : m_wDiskName(wDiskName),
																m_eDevType(m_eDevType),
																m_eIoType(eIoType)
{
}

Disk::~Disk(void)
{
}


bool Disk::OpenDisk()
{
	// Open the disk
	SetErrorMode(SEM_FAILCRITICALERRORS);

	if (m_eIoType == eWrite)
	{
		m_phDiskFile = CreateFile(m_wDiskName.c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ |
			       FILE_SHARE_WRITE, NULL, OPEN_ALWAYS,
			       FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED| FILE_FLAG_WRITE_THROUGH, NULL);
	}
	else
	{
		// Reading file just if it exists
		m_phDiskFile = CreateFile(m_wDiskName.c_str(), GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ |
		       FILE_SHARE_WRITE, NULL, OPEN_EXISTING,
		       FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED, NULL);
	}

	if (m_phDiskFile == INVALID_HANDLE_VALUE)
	{
		cout << "failed in creating file - error: " << GetLastError()<< endl;
		return false;
	}
	else
	{
		if (m_eDevType == ePhysical)
		{
			// Lock the disk and unmount it for physical writing
			
			// Lock the device

			DWORD  lpBytesReturned = 0;
			BOOL bRet = DeviceIoControl(m_phDiskFile,  
						FSCTL_LOCK_VOLUME,
						NULL, 0, // no input buffer
						NULL, 0, // no input buffer
						&lpBytesReturned,// # bytes returned
						NULL);  // asynchronous I/


			if(!bRet)
			{
				cout << "*** Error " << GetLastError()<< " closing disk" << endl;
				CloseHandle(m_phDiskFile);
				return false;
			}

#ifdef __FLASH_DEBUG__
			cout << "Locked the device" << endl;
#endif // __FLASH_DEBUG__

			// Dismount the disk
			bRet = DeviceIoControl(
				(HANDLE) m_phDiskFile,            // handle to a volume
				(DWORD) FSCTL_DISMOUNT_VOLUME,   // dwIoControlCode
				NULL,                        // lpInBuffer
				0,                           // nInBufferSize
				NULL,                        // lpOutBuffer
				0,                           // nOutBufferSize
				&lpBytesReturned,   // number of bytes returned
				NULL
				);

			if(!bRet)
			{
				cout << "*** Error " << GetLastError()<< " closing disk" << endl;
				CloseHandle(m_phDiskFile);
				return false;
			}

#ifdef __FLASH_DEBUG__
			cout << "dismounted the device" << endl;
#endif // __FLASH_DEBUG__

			//// Get rid of dos device

			//TCHAR szDriveLetter[3];
			//szDriveLetter[1] = TEXT('E');
			//szDriveLetter[1] = TEXT(':');
			//szDriveLetter[2] = TEXT('\0');

			//bRet = DefineDosDevice (DDD_RAW_TARGET_PATH|DDD_REMOVE_DEFINITION | DDD_EXACT_MATCH_ON_REMOVE,
			//	szDriveLetter, NULL);

			//if(!bRet)
			//{
			//	cout << "*** Error " << GetLastError()<< " get rid of dos device" << endl;
			//	CloseHandle(m_phDiskFile);
			//	return false;
			//}

		}

		if (CreateIoCompletionPort(m_phDiskFile, m_hCompletionQueue, 0, 1)== NULL)
		{
			// TODO: add here staff if needed
		}
		else
		{	
			// TODO: add here staff if needed
		}
	}
	return true;

}

bool Disk::CloseDisk()
{
	bool bRes = false;

	// Close the event
	CloseHandle(m_oOverlapInfo.hEvent);
	
	// UnLock the physical disk
	if (m_eDevType == ePhysical)
	{
		DWORD dwBytesReturned = 0;
		
		BOOL bRet = DeviceIoControl(m_phDiskFile,  
			FSCTL_UNLOCK_VOLUME,
			NULL, 0, // no input buffer
			NULL, 0, // no input buffer
			&dwBytesReturned,// # bytes returned
			NULL);  // synchronous I/

		if (bRet != TRUE)
		{
#ifdef __FLASH_DEBUG__
			cout << "*** Error " << GetLastError()<< " unlocking disk" << endl;
#endif // __FLASH_DEBUG__
			bRes =  false;
		}
	}

	// Close the disk
	if (!CloseHandle(m_phDiskFile))
	{
		cout << "*** Error " << GetLastError()<< " closing disk" << endl;
		SetLastError(0);
		bRes = false;
	}
	else
	{
#ifdef __FLASH_DEBUG__
		cout << "Succeeded in closing the disk" << endl;
#endif // __FLASH_DEBUG__
		bRes = true;
	}



	return bRes;

}

bool Disk::CreateOverlappedEvent()
{
	bool bIsReady = true;
	// Create an event.
	m_oOverlapInfo.hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);

	// Was the event created successfully?
	if (m_oOverlapInfo.hEvent == NULL) {
		cout << "*** CreateEvent() failed in CTestSegment::PrepareTest()" << endl;
		bIsReady = false;
	}
	else
	{
		// Set the low-order bit of the event handle to prevent it from being
		// picked up by a later call to GetQueuedCompletionStatus() 
		m_oOverlapInfo.hEvent = (HANDLE) ((UINT_PTR) m_oOverlapInfo.hEvent | 0x0000000000000001);

		// Set its address. the offset is inside the file - for example if you want offset 10 
		// so the written file size will be the offset + written buffer size

	}
	return bIsReady;
}

bool Disk::DoOverlappedIO(void* pBuffer, int nBytesCount, DWORDLONG dwlPrepareOffset /*= 0*/, int* pnOpSectors /*=NULL*/, float* pfIoTimeMsec /*= NULL*/)
{
	// Take the start time in order to check the IO time
	LARGE_INTEGER  startTime;
	QueryPerformanceCounter((LARGE_INTEGER *)&startTime);
	DWORD dwStart = timeGetTime();

	DWORD dwBytesWritten = 0;
	BOOL bRes = FALSE;
	string sOp;

	m_oOverlapInfo.Offset		= (DWORD)dwlPrepareOffset;
	m_oOverlapInfo.OffsetHigh	= (DWORD)(dwlPrepareOffset >> 32);

	if (m_eIoType == eWrite)
	{
		bRes = WriteFile(m_phDiskFile, (char *)pBuffer, 
			nBytesCount, &dwBytesWritten, &(m_oOverlapInfo));
#ifdef __FLASH_DEBUG__
		sOp = "writing";
#endif // #ifdef __FLASH_DEBUG__
	}
	else
	{
		bRes = ReadFile(m_phDiskFile, (char *)pBuffer, 
			nBytesCount, &dwBytesWritten, &(m_oOverlapInfo));
#ifdef __FLASH_DEBUG__
		sOp = "reading";
#endif // __FLASH_DEBUG__
	}

	// Do the asynchronous IO.
	if (bRes == TRUE) {
			// It succeeded immediately!
#ifdef __FLASH_DEBUG__
			cout << sOp << "(immediately) " << dwBytesWritten << " of "
			    << nBytesCount << " bytes to disk ";
			std::wcout << m_wDiskName << "." << endl;
#endif // __FLASH_DEBUG__
	}
	else {
			if (GetLastError() != ERROR_IO_PENDING) 
			{ 
				 // WriteFile / ReadFile failed, but isn't delayed. Report error and abort.
				// It didn't start OK!  Bail out.
				cout << "***Error (immediately) " << sOp << dwBytesWritten <<
					" of " << nBytesCount << " bytes to disk: ";
				std::wcout << m_wDiskName;
				cout <<", error=" << GetLastError() << endl;
				CloseHandle(m_oOverlapInfo.hEvent);
				return false;
			}
	}
	 // operation is pending.
     DWORD dwRes = WaitForSingleObject(m_oOverlapInfo.hEvent, INFINITE);
     switch(dwRes)
     {
        // OVERLAPPED structure's event has been signaled. 
        case WAIT_OBJECT_0:
		{

			// Check to see if it has completed
			if (GetOverlappedResult(m_phDiskFile, &(m_oOverlapInfo), &dwBytesWritten, FALSE)) 
			{
				// It completed successfully!
#ifdef __FLASH_DEBUG__
				cout << sOp << " (eventually) " << dwBytesWritten << " of " << nBytesCount
					<< ", offset = "<< m_oOverlapInfo.Offset << " bytes to disk ";
					std::wcout << m_wDiskName << "." << endl;
#endif // __FLASH_DEBUG__
			}
			else
			{
				// It  wasn't completed successfully!
				cout << "***Error (eventually) " << sOp  << dwBytesWritten << " of " << nBytesCount
					<< " bytes to disk ";
				std::wcout << m_wDiskName;
				cout << "."   << "error=" << GetLastError() << endl;
				CloseHandle(m_oOverlapInfo.hEvent);
				return false;
			}
		}

	 }

	 if (NULL != pnOpSectors)
	 {
		*pnOpSectors += (dwBytesWritten/512);
//		printf("operation on %d sectors \n", *pnOpSectors);
	 }

	LARGE_INTEGER  endTime;
	QueryPerformanceCounter((LARGE_INTEGER *)&endTime);

	LARGE_INTEGER freq;
	QueryPerformanceFrequency((LARGE_INTEGER*)&freq);
	
	double timeDifferenceInMilliseconds = ((endTime.QuadPart -startTime.QuadPart) * 1000.0 / freq.QuadPart);

	if (NULL != pfIoTimeMsec)
	{
		*pfIoTimeMsec += static_cast<float>(timeDifferenceInMilliseconds);
	}

	return true;
}
